import json
import os.path
from datetime import datetime, date

from flask import Flask, render_template, request
from flask.json import JSONEncoder, jsonify
from flask_sqlalchemy import SQLAlchemy
from mlwrite.config import BASE_DIR, SQLALCHEMY_DATABASE_URI, RedisConfig, MODEL_DICT, DEVICE
from redis import StrictRedis

db = SQLAlchemy()
redis_template = StrictRedis(host=RedisConfig.REDIS_HOST, port=int(RedisConfig.REDIS_PORT))


def create_app(test_config=None):
    # instance_relative_config 开启允许从文件中加载配置
    app = Flask(__name__, instance_relative_config=True)

    if test_config is None:
        # load the instance config, if it exists, when not testing
        CONFIG_PATH = BASE_DIR / 'mlwrite/config.py'
        app.config.from_pyfile(CONFIG_PATH, silent=True)
    else:
        # load the test config if passed in
        app.config.from_mapping(test_config)

    # 数据库连接对象
    db.init_app(app)

    # 替换json编码器，格式化时间
    app.json_encoder = CustomJSONEncoder

    # ensure the instance folder exists
    try:
        os.makedirs(app.instance_path)
    except OSError:
        pass

    # 引入蓝图
    from app.write import views as write
    app.register_blueprint(write.bp)

    from app.admin import views as admin
    app.register_blueprint(admin.bp)

    from app.train import views as train
    app.register_blueprint(train.bp)

    # 定义首页
    app.add_url_rule('/', endpoint='index')

    # 加载模型
    with app.app_context():
        load_model()

    # print(MODEL_DICT)

    # a simple page that says hello
    @app.route('/')
    @app.route('/index.html')
    def index():
        return render_template(template_name_or_list='index.html')

        # 静态页面统一跳转

    @app.route('/page/<path:page>')
    def page_router(page):
        args = json.dumps(request.args.to_dict())
        return render_template('/page/' + page, args=args)

    return app


# 日期格式化
class CustomJSONEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime('%Y-%m-%d %H:%M:%S')
        elif isinstance(obj, date):
            return obj.strftime('%Y-%m-%d')
        else:
            return JSONEncoder.default(self, obj)


def load_model():
    import torch
    import app.admin.models
    from app.admin.models import TModules
    from app.write.model import GPT, GPTConfig
    from torch.nn import functional as F
    models = TModules.query.filter(app.admin.models.TModules.module_path != '')\
        .filter(app.admin.models.TModules.init_load == 1)\
        .all()
    for module in models:
        model_name = module.module_path
        word_name = module.vocab_path
        ctx_len = module.ctx_len  # 上下文长度
        n_layer = 12  # 隐藏层
        n_head = 12
        n_embd = n_head * 64
        n_attn = n_embd
        n_ffn = n_embd
        top_p = 0.75  # 这个的范围是 0 到 1。越大，变化越多。越小，生成效果越规矩。自己试试 0 和 0.5 和 1.0 的效果就知道了
        top_p_newline = 0.9
        # 打开模型对应词库加载为词典
        with open(word_name, "r", encoding="utf-16") as result_file:
            word_table = json.load(result_file)
        # 词典长度
        vocab_size = len(word_table)
        train_dataset = lambda: None

        # 词典映射
        train_dataset.stoi = {v: int(k) for k, v in word_table.items()}
        train_dataset.itos = {int(k): v for k, v in word_table.items()}
        # 未知字符替换
        unknown_char = train_dataset.stoi['.']
        model = GPT(
            GPTConfig(vocab_size, ctx_len, n_layer=n_layer, n_head=n_head, n_embd=n_embd, n_attn=n_attn,
                      n_ffn=n_ffn))
        m2 = torch.load(model_name, map_location='cpu').state_dict()
        for i in range(n_layer):
            prefix = f'blocks.{i}.attn.'
            time_w = m2[prefix + 'time_w']
            time_alpha = m2[prefix + 'time_alpha']
            time_beta = m2[prefix + 'time_beta']

            TT = ctx_len
            T = ctx_len
            w = F.pad(time_w, (0, TT))
            w = torch.tile(w, [TT])
            w = w[:, :-TT].reshape(-1, TT, 2 * TT - 1)
            w = w[:, :, TT - 1:]
            w = w[:, :T, :T] * time_alpha[:, :, :T] * time_beta[:, :T, :]

            m2[prefix + 'time_ww'] = w
            del m2[prefix + 'time_w']
            del m2[prefix + 'time_alpha']
            del m2[prefix + 'time_beta']
        if DEVICE == 'gpu':
            if torch.cuda.is_available():
                model = model.cuda()
        model.load_state_dict(m2)
        model_dict = {
            "train_dataset": train_dataset,
            "model": model,
            "m2": m2,
            "unknown_char": unknown_char,
            "ctx_len": ctx_len,
            "device": DEVICE,
            "top_p": top_p,
            "top_p_newline": top_p_newline
        }
        MODEL_DICT[module.module_code] = model_dict
